#!/usr/bin/python
""" This program filters reads by expression or file. """
import os
import sys

def search(name):
	for f in filters:
		if name.count(f) > 0: return f
	return False

# Usage.
if len(sys.argv) < 2:
	print ""
	print "Usage: %s [-p id | -f file_list | -l 45 ] [-c start stop] -in fasta_file" % sys.argv[0]
	print "-p: if present will perform filter in by id"
	print "-f: filters by id's in a file, one per line."
	print "-l: filters by sequence length, GTE"
	print "-c: filters by relative start and stop (inclusive). In addition to required filter."
	print "-fc: filters by relative start and stop included in file list e.g. 'name'{\\t}start{\\t}stop"
	print "-in: input fasta file."
	print "-inc: only print seqs which match list"
	print "-exc: default, print seqs which dont match"
	print "-inexact: searches headers for id, header might have more info. takes longer."
	print ""
	sys.exit()

# Parse argse
inc = False
dashf = False
dashc = False
filterbylen = False
inexact = False
for i in range(1,len(sys.argv)):
	if sys.argv[i] == "-p":
		fid = sys.argv[i+1]
	if sys.argv[i] == "-f":
		flist = sys.argv[i+1]
		dashf = True
	if sys.argv[i] == "-in":
		fastain = sys.argv[i+1]
	if sys.argv[i] == "-inc":
		inc = True	
	if sys.argv[i] == "-inexact":
		inexact = True	
	if sys.argv[i] == "-fc":
		dashc = True
	if sys.argv[i] == "-c":
		dashc = True
		dc_start = int(sys.argv[i+1])
		dc_stop = int(sys.argv[i+2])
	if sys.argv[i] == "-l":
		filterbylen = True
		filtlen = int( sys.argv[i+1] )

# Setup filter list.
filters = {}
if filterbylen == False:
	if dashf == False:
		filters[fid] = True
	else:
		# length specified.
		if dashc == True:
			fin = open(flist,"r")
			for f in fin:
				if f.count("\t") > 0:
					tmp = f.strip().split("\t")	
				elif f.count(" ") > 0:
					tmp = f.strip().split(" ")
				else:
					print "Unknown filter file format"
					sys.exit()			
				filters[tmp[0].replace(">","")] = [int(tmp[1]), int(tmp[2])]
			fin.close()
		else:
			fin = open(flist,"r")
			for f in fin:
				filters[f.replace("\n","").replace(">","")] = True
			fin.close()



# Check each sequence against filter.
seq = ""
fin = open(fastain,"r")
for line in fin:
	# Skip blanks.
	if len(line) < 2: continue
	if line[0] == "#": continue

	# Check for ids.
	if line.count(">") > 0: 

		# Check if ending seq.
		if len(seq) > 0:
			tmp = head.replace(">","").strip()
			toprint = False

			# If filter in.
			if inc == True:
				if filterbylen == True and len(seq) >= filtlen: toprint = True	
				elif tmp in filters:
					if dashc == True:
						dc_start = filters[tmp][0]
						dc_stop = filters[tmp][1]
					toprint = True			
				elif inexact == True:
					res = search(tmp)
					if res != False: 
						if dashc == True:
							dc_start = filters[res][0]
							dc_stop = filters[res][1]
						toprint = True
					else: toprint = False
					
			else:	
				if filterbylen == True and len(seq) < filtlen: toprint = True
				if tmp not in filters: toprint = True
				elif inexact == True:
					if search(tmp) == True: toprint = False
					else: toprint = True

			# Print if appropriate.
			if toprint == True:
				if dashc == True:
					# Print subseq.
					print head.replace("\n","")
					print seq[dc_start:dc_stop].replace("\n","")
				else:
					# Print full seq.
					print head.replace("\n","")
					print seq.replace("\n","")

			# Setup next.
			head = line
			seq = ""
			continue

		# Initial seq.
		head = line
		continue

	# Filter chars.
	seq += line

# If filter in.
tmp = head.replace(">","").strip()
toprint = False
if inc == True:
	if filterbylen == True and len(seq) >= filtlen: toprint = True	
	elif tmp in filters:
		if dashc == True:
			dc_start = filters[tmp][0]
			dc_stop = filters[tmp][1]
		toprint = True			
	elif inexact == True:
		res = search(tmp)
		if res != False: 
			if dashc == True:
				dc_start = filters[res][0]
				dc_stop = filters[res][1]
			toprint = True
		else: toprint = False
		
else:	
	if filterbylen == True and len(seq) < filtlen: toprint = True
	if tmp not in filters: toprint = True
	elif inexact == True:
		if search(tmp) == True: toprint = False
		else: toprint = True


# Close outifle
fin.close()



